import argparse
import gc

import torch
from torch.utils.data import DataLoader
import torch.optim as optim
from Models.Yolo import YOLO
from Models.YoloLoss import YOLOLoss
from Utils import ModelUtils
from Config.ConfigTrain import *
from Utils.DataSet.VOC import VOCDataSet
import os
from Utils import SaveModel
from Utils import Log
from torch.utils.tensorboard import SummaryWriter


def train():

    ModelUtils.set_seed()
    # 初始化驱动
    device = None
    if (torch.cuda.is_available()):
        if (not opt.device == 'cpu'):
            div = "cuda:" + opt.device
            device = torch.device(div)
            torch.backends.cudnn.benchmark = True
            print("\033[0;31;0m使用GPU训练中：{}\033[0m".format(torch.cuda.get_device_name()))
        else:
            device = torch.device("cpu")
            print("\033[0;31;40m使用CPU训练\033[0m")
    else:
        device = torch.device("cpu")
        print("\033[0;31;40m使用CPU训练\033[0m")

    # 创建 runs exp 文件
    EPX_Path = SaveModel.CreatRun(0,"detect")
    # 日志相关的准备工作
    wirter = None
    openTensorboard = opt.tensorboardopen
    path_board = None
    if (openTensorboard):
        path_board = EPX_Path + "\\logs"
        wirter = SummaryWriter(path_board)

    fo = Log.PrintLog(EPX_Path)


    train_data_dir_image = opt.train_dir_image
    train_data_dir_Ann = opt.train_dir_Ann
    if (not train_data_dir_image):
        train_data_dir_image = TrainImage
        if (not os.path.exists(train_data_dir_image)):
            raise Exception("训练集路径错误")
    if (not train_data_dir_Ann):
        train_data_dir_Ann = TrainAnn
        if (not os.path.exists(train_data_dir_Ann)):
            raise Exception("训练集路径错误")



    train_data =VOCDataSet(imgs_path=train_data_dir_image,
                               annotations_path=train_data_dir_Ann,
                               is_train=True)

    valid_data_dir_image = opt.valid_dir_image
    valid_data_dir_Ann = opt.valid_dir_Ann

    if (not valid_data_dir_image):
        valid_data_dir_image = ValImage
        if (not os.path.exists(valid_data_dir_image)):
            raise Exception("训练集路径错误")
    if (not valid_data_dir_Ann):
        valid_data_dir_Ann = ValAnn
        if (not os.path.exists(valid_data_dir_Ann)):
            raise Exception("训练集路径错误")

    valid_data = VOCDataSet(imgs_path=valid_data_dir_image,
                             annotations_path=valid_data_dir_Ann,
                             is_train=False)

    # 构建DataLoder
    train_loader = DataLoader(dataset=train_data, batch_size=opt.batch_size, num_workers=opt.works, shuffle=True)
    valid_loader = DataLoader(dataset=valid_data, batch_size=opt.batch_size)
    # 1 开始初始化网络，设置参数啥的
    net = YOLO(B=2,classes_num=Classes)
    #加载预训练权重
    if(PreWeight):
        # 1.1 初始化网络
        preweight = torch.load(PreWeight)
        net.initialize_weights(preweight)

    net = net.to(device)


    loss_func = YOLOLoss(S=7,B=2,Classes=Classes).to(device)

    # 1.3设置优化器
    optimizer = optim.SGD(net.parameters(), lr=opt.lr, momentum=0.09)  # 选择优化器
    # 设置学习率下降策略,默认的也可以，那就不设置嘛，主要是不断去自动调整学习的那个速度
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.01)

    # 2 开始进入训练步骤
    # 2.1 进入网络训练

    Best_weight = None

    TotalLoss = 0.
    ValLoss = 0.
    ValTime = 0.
    Best_loss = float("inf")
    for epoch in range(opt.epochs):

        """
        下面是一些用来记录当前网络运行状态的参数
        """
        train_loss = 0
        val_loss = 0
        # train_iou = 0
        # val_iou = 0
        # train_object_num = 0
        # val_object_num = 0
        train_loss_coord = 0
        val_loss_coord = 0
        train_loss_confidence = 0
        val_loss_confidence = 0
        train_loss_classes = 0
        val_loss_classes = 0

        log_loss_mean_train = 0.
        # log_loss_mean_val = 0.
        net.train()
        print("正在进行第{}轮训练".format(epoch + 1))
        for i, data in enumerate(train_loader):

            # forward
            inputs, labels = data
            inputs, labels = inputs.float().to(device), labels.float().to(device)
            outputs = net(inputs)
            optimizer.zero_grad()

            loss = loss_func(bounding_boxes=outputs, ground_truth=labels,batch_size = opt.batch_size )
            batch_loss = loss[0]
            batch_loss.backward()
            optimizer.step()
            log_loss_mean_train+=batch_loss
            train_loss+=batch_loss
            train_loss_coord+=loss[1]
            train_loss_confidence+=loss[2]
            train_loss_classes+=loss[3]

            # train_iou+=train_iou+loss[4]
            # train_object_num+=loss[5]

            # update weights


            if (i + 1) % opt.log_interval == 0:
                log_loss_mean_train = log_loss_mean_train / opt.log_interval

                info = "训练:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f}" \
                    .format \
                        (
                        epoch, opt.epochs, i + 1, len(train_loader), log_loss_mean_train
                    )
                print(info, file=fo)

                if (opt.show_log_console):
                    info_print = "\033[0;33;0m" + info + "\033[0m"
                    print(info_print)

                log_loss_mean_train = 0.0
        #总体损失
        TotalLoss+=train_loss
        # tensorboard 绘图
        if (wirter):

            wirter.add_scalar("总体损失值",TotalLoss,epoch)
            wirter.add_scalar("每轮损失值",train_loss,epoch)
            wirter.add_scalar("每轮预测预测框损失值",train_loss_coord,epoch)
            wirter.add_scalar("每轮预测框置信度损失",train_loss_confidence,epoch)
            wirter.add_scalar("每轮预测类别损失值",train_loss_classes,epoch)


        # 保存效果最好的玩意
        if (train_loss < Best_loss):
            Best_weight = net.state_dict()
            Best_loss = train_loss

        scheduler.step()  # 更新学习率

        # 2.2 进入训练对比阶段

        if (epoch + 1) % opt.val_interval == 0:

            """
            这部分和训练的那部分是类似的，可以忽略这部分的代码
            """

            net.eval()
            with torch.no_grad():
                for j, data in enumerate(valid_loader):
                    inputs, labels = data
                    inputs, labels = inputs.float().to(device), labels.float().to(device)
                    outputs = net(inputs)
                    loss = loss_func(outputs, labels)

                    batch_loss = loss[0]
                    # log_loss_mean_val += batchLoss
                    val_loss += batch_loss
                    val_loss_coord += loss[1]
                    val_loss_confidence += loss[2]
                    val_loss_classes += loss[3]
                    # val_iou += train_iou + loss[4]
                    # val_object_num += loss[5]

                info_val = "测试:\tEpoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} ".format \
                        (
                        epoch, opt.epochs, (j+1), len(valid_loader), val_loss
                    )
                print(info_val, file=fo)
                if (opt.show_log_console):
                    info_print_val = "\033[0;31;0m" + info_val + "\033[0m"
                    print(info_print_val)
                ValLoss+=val_loss
                if (wirter):
                    wirter.add_scalar("测试总体损失",ValLoss, (ValTime))
                    wirter.add_scalar("每次测试总损失总值", val_loss, (ValTime))
                    wirter.add_scalar("每轮测试预测框损失值", val_loss_coord, ValTime)
                    wirter.add_scalar("每轮测试预测框置信度损失", val_loss_confidence, ValTime)
                    wirter.add_scalar("每轮测试预测类别损失值", val_loss_classes, ValTime)
                ValTime+=1



    # 最后一次的权重
    Last_weight = net.state_dict()

    # 保存模型
    SaveModel.Save_Model(EPX_Path, Best_weight, Last_weight)

    fo.close()
    if (wirter):
        print("tensorboard dir is:", path_board)
    wirter.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=300)
    parser.add_argument('--batch-size', type=int, default=4)
    parser.add_argument('--lr', type=float, default=0.01)
    #每5个batch输出一次结果
    parser.add_argument('--log_interval', type=int, default=2)
    # 训练几轮测试一次
    parser.add_argument('--val_interval', type=int, default=10)
    parser.add_argument('--train_dir_image', type=str, default='')
    parser.add_argument('--train_dir_Ann', type=str, default='')
    parser.add_argument('--valid_dir_image', type=str, default='')
    parser.add_argument('--valid_dir_Ann', type=str, default='')
    # 如果是Mac系注意这个参数可能需要设置为1，本地训练，不推荐MAC
    parser.add_argument('--works', type=int, default=0)
    parser.add_argument('--show_log_console', type=bool, default=True)
    parser.add_argument('--device', type=str, default="cpu", help="默认使用显卡加速训练参数选择：0,1,2...or cpu")
    parser.add_argument('--tensorboardopen', type=bool, default=True)
    opt = parser.parse_args()
    train()

    # tensorboard --logdir = runs/train/epx2/logs

